package com.webcat.service.impl;

import com.alibaba.fastjson.JSON;
import com.webcat.constants.PromptTemplateConstant;
import com.webcat.dify.kb.KBResponse;
import com.webcat.dify.kb.entity.Record;
import com.webcat.entity.SseResponse;
import com.webcat.entity.vo.TopicVo;
import com.webcat.service.AbstractChatService;
import com.webcat.service.QuestionService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class QuestionServiceImpl extends AbstractChatService implements QuestionService {

    @Override
    public String start() {
        return super.start();
    }

    @Override
    public Flux<SseResponse<Object>> streamChat(String chatId, String message, String datasetId, String questionType, Boolean multiConversation) {
        if (datasetId != null && !datasetId.isEmpty()) {
            KBResponse kbResponse = difyKBClient.getSyncKbResponse(message, datasetId);
            List<Record> records  = kbResponse.getRecords();
            String ragContent = generateRagContent(records);
            PromptTemplate template = new PromptTemplate(PromptTemplateConstant.RAG_PROMPT_TEMPLATE);
            UserMessage userMessage = (UserMessage)template.createMessage(Map.of("question", message, "context", ragContent));
            Prompt prompt = new Prompt(userMessage);
            Flux<ChatResponse> chatResponseFlux = this.generateChatClient(chatId, multiConversation, false).prompt(prompt).stream().chatResponse();
            return chatResponseFlux.map(response -> {
                String finishReason = response.getResult().getMetadata().getFinishReason();
                if ("stop".equalsIgnoreCase(finishReason)) {
                    //return SseResponse.success(chatId, "end", generateKnowledgeBase(datasetId, records));
                    return SseResponse.success(chatId, "end", "");
                }
                else {
                    String text = response.getResult().getOutput().getText();
                    return SseResponse.success(chatId, "message", text);
                }
            });
        }
        else {
            Prompt prompt = new Prompt(message);
            Flux<ChatResponse> chatResponseFlux = this.generateChatClient(chatId, multiConversation ,true).prompt(prompt).stream().chatResponse();
            return chatResponseFlux.map(response -> generateResponse(chatId, response));
        }
    }

    @Override
    public String chat(String chatId, String message, String datasetId, String questionType, Boolean multiConversation) {
        return null;
    }

    @Override
    public List<TopicVo> editQuestion(String chatId, String message) {
        PromptTemplate template = new PromptTemplate(PromptTemplateConstant.EDIT_QUESTION_JSON_PROMPT_TEMPLATE);
        Prompt prompt = new Prompt(template.createMessage(Map.of("context", message)));
        ChatResponse response = this.generateChatClient(chatId, false ,true).prompt(prompt).call().chatResponse();
        assert response != null;
        String result = response.getResult().getOutput().getText();
        result = result.replace("```json", "");
        result = result.replace("```", "");
        return JSON.parseArray(result, TopicVo.class);
    }

    private Flux<SseResponse<Object>> generateQuestionResult(String chatId, String datasetId, boolean multiConversation, Prompt prompt, List<Record> records) {
        ChatResponse chatResponse = this.generateChatClient(chatId, multiConversation, false).prompt(prompt).call().chatResponse();
            assert chatResponse != null;
            String result = chatResponse.getResult().getOutput().getText();
            log.info("result={}", result);
            result = result.replace("```json", "");
            result = result.replace("```", "");
            List<TopicVo> topicList = JSON.parseArray(result, TopicVo.class);
            if (topicList == null || topicList.isEmpty()) {
                return Flux.just(SseResponse.success(chatId, "message", "对不起，我不知道"));
            }

            List<SseResponse<Object>> sseResponseList = new ArrayList<>(topicList.stream().map(topic -> {
                StringBuilder text = new StringBuilder();
                text.append(topic.getQuestion()).append("\n");
                List<String> options = topic.getOptions();
                if (options != null && !options.isEmpty()) {
                    options.forEach(op -> text.append(op).append("\n"));
                }
                text.append("答案:").append(topic.getAnswer()).append("\n");
                if (topic.getReference() != null && !topic.getReference().isEmpty()) {
                    text.append("引用:").append(topic.getReference()).append("\n");
                }
                return SseResponse.success(chatId, "message", (Object) text.toString());
            }).toList());
            List<SseResponse<Object>> jsonSseResponseList = new ArrayList<>(topicList.stream().map(topic -> {
                return SseResponse.success(chatId, "jsonMessage", (Object) topic);
            }).toList());

            sseResponseList.addAll(jsonSseResponseList);

            /*KnowledgeBaseVo knowledgeBaseVo;
            if (records != null && !records.isEmpty()) {
                knowledgeBaseVo = new KnowledgeBaseVo();
                knowledgeBaseVo.setDatasetId(datasetId);
                List<DocumentVo> documentVoList = records.stream().map(record -> {
                    DocumentVo documentVo = new DocumentVo();
                    documentVo.setId(record.getSegment().getDocument().getId());
                    documentVo.setName(record.getSegment().getDocument().getName());
                    documentVo.setContent(record.getSegment().getContent());
                    return documentVo;
                }).toList();
                knowledgeBaseVo.setDocumentList(documentVoList);
            }*/
            sseResponseList.add(SseResponse.success(chatId, "end", ""));
            //sseResponseList.add(SseResponse.success(chatId, "end", generateKnowledgeBase(datasetId, records)));
            return Flux.fromIterable(sseResponseList);
    }
}
